import torch
from torch.utils.data._utils.collate import default_collate

def merge(x, y):
    if isinstance(x, torch.Tensor):
        return torch.cat((x, y), dim=0)
    elif isinstance(x, list) or isinstance(x, tuple):
        return x + y
    else:
        raise NotImplementedError("merge not implemented for type {}".format(type(x)))

def collate_fn(batch, self_len=None):
    if not sum(self_len):
        return default_collate(batch)
    elif len(batch) > sum(self_len):
        # collate twice for in-domain and out-domain part
        self_len = sum(self_len)    # to be fixed, this is a simple case
        in_domain_batch = batch[:-self_len]
        out_domain_batch = batch[-self_len:]
        in_domain_batch = default_collate(in_domain_batch)
        out_domain_batch = default_collate(out_domain_batch)
        # merge element after first element
        assert len(in_domain_batch) <= len(out_domain_batch), "wrong setup with two datasets"    # out_domain_batch contains the crop lenth
        other_elements = [merge(x, y) for x, y in zip(in_domain_batch[1:], out_domain_batch[1:])]
        crop_pos = out_domain_batch[-1]
        return in_domain_batch[0], out_domain_batch[0], *other_elements, crop_pos
    else:
        return default_collate(batch)
